package scu.maqiang.numeric;

import scu.maqiang.mesh.ScalarFunc;
import scu.maqiang.mesh.Tecplot;
import scu.maqiang.mesh.VectorFunc;

import java.util.Arrays;
import java.util.Vector;
import java.util.function.Function;

public class ODESolver {

	public static Pair<double[], double[]> Euler(double[] inter, ScalarFunc func, double y0, int n) {
		double[] t = new double[n+1];
		double[] y = new double[n+1];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[] time = {0};
		double[] yc = {0};
		for(int i = 0; i < n; i++) {
			time[0] = t[i];
			yc[0] = y[i];
			t[i+1] = t[i] + h;
			y[i+1] = y[i] + h * func.action(yc, 0, time);
		}
		return new Pair<>(t, y);
	}

	public static Pair<double[], double[][]> Euler(double[] inter, VectorFunc func, double y0[], int n) {
		double[] t = new double[n+1];
		int ny = y0.length;
		double[][] y = new double[n+1][ny];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[][] time = {{0}};
		for(int i = 0; i < n; i++) {
			time[0][0] = t[i];
			double[] yi = y[i];
			t[i+1] = t[i] + h;
			double[] funVal = func.action(yi, 0, time);
			Arrays.setAll(y[i + 1], j -> yi[j] + h * funVal[j]);
		}
		return new Pair<>(t, y);
	}

	public static Pair<double[], double[]> trapezoidalMethod(double[] inter, ScalarFunc func, double y0, int n) {
		double[] t = new double[n+1];
		double[] y = new double[n+1];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[] time = {0};
		double[] yc = {0};
		for(int i = 0; i < n; i++) {
			time[0] = t[i];
			yc[0] = y[i];
			t[i+1] = t[i] + h;
			double func1 = func.action(yc, 0, time);
			double temp = y[i] + h * func1;
			time[0] = t[i + 1];
			yc[0] = temp;
			double func2 = func.action(yc, 0, time);
			y[i+1] = y[i] + 0.5 * h * (func1 + func2);
		}
		return new Pair<>(t, y);
	}

	public static Pair<double[], double[][]> trapezoidalMethod(double[] inter, VectorFunc func, double y0[], int n) {
		double[] t = new double[n+1];
		int ny = y0.length;
		double[][] y = new double[n+1][ny];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[][] time = {{0}};
		double[] temp = new double[ny];
		for(int i = 0; i < n; i++) {
			time[0][0] = t[i];
			double[] yi = y[i];
			t[i+1] = t[i] + h;
			double[] fun1 = func.action(yi, 0, time);
			Arrays.setAll(temp, j -> yi[j] + h * fun1[j]);
			time[0][0] = t[i + 1];
			double[] fun2 = func.action(temp, 0, time);
			Arrays.setAll(y[i + 1], j -> yi[j] + 0.5 * h * (fun1[j] + fun2[j]));
		}
		return new Pair<>(t, y);
	}
	
	public static Pair<double[], double[]> RungeKutta44(double[] inter, Function<double[], Double> func, double y0, int n) {
		double[] t = new double[n+1];
		double[] y = new double[n+1];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[] ty = new double[2];
		for(int i = 0; i < n; i++) {
			t[i+1] = t[i] + h;
			ty[0] = t[i];
			ty[1] = y[i];
			double s1 = func.apply(ty);
			ty[0] += 0.5 * h;
			ty[1] += 0.5 * h * s1;
			double s2 = func.apply(ty);
			ty[1] = y[i] + 0.5 * h * s2;
			double s3 = func.apply(ty);
			ty[0] += 0.5 * h;
			ty[1] = y[i] + h * s3;
			double s4 = func.apply(ty);
			y[i+1] = y[i] + h * (s1 + 2 * s2 + 2 * s3 + s4) / 6;
		}
		return new Pair<>(t, y);
	}



	public static Pair<double[], double[][]> RungeKutta44(double[] inter, Function<double[], double[]> func, double[] y0, int n) {
		double[] t = new double[n+1];
		int ny = y0.length;
		double[][] y = new double[n+1][ny];
		double h = (inter[1] - inter[0]) / n;
		t[0] = inter[0];
		y[0] = y0;
		double[] ty = new double[ny + 1];
		for(int i = 0; i < n; i++) {
			t[i+1] = t[i] + h;
			ty[0] = t[i];
			double[] s1 = func.apply(ty);
			for(int j = 1; j <= ny; j++) {
				ty[j] = y[i][j - 1];
			}
			ty[0] += 0.5 * h;
			for(int j = 1; j <= ny; j++) {
				ty[j] += 0.5 * h * s1[j - 1];
			}
			double[] s2 = func.apply(ty);
			for(int j = 1; j <= ny; j++) {
				ty[j] = y[i][j-1] + 0.5 * h * s2[j-1];
			}
			double[] s3 = func.apply(ty);
			ty[0] += 0.5 * h;
			for(int j = 1; j <= ny; j++) {
				ty[j] = y[i][j-1] + h * s3[j-1];
			}
			double[] s4 = func.apply(ty);
			for(int j = 0; j < ny; j++) {
				y[i+1][j] = y[i][j] + h * (s1[j] + 2 * s2[j] + 2 * s3[j] + s4[j]) / 6;
			}
		}
		return new Pair<>(t, y);
	}
	
	public static void main(String[] args) {
//		Function<double[], Double> func = ty -> ty[0] * ty[1] + ty[0] * ty[0] * ty[0];
//		Pair<double[], double[]> result = Euler(new double[] {0, 1}, func, 1, 5);
//		System.out.println(Arrays.toString(result.getSecond()));
//		MVO.toFile(result.getFirst(), result.getSecond(), "result.dat");

		//ScalarFunc f = (y, label, t) -> t[0] * y[0] + t[0] * t[0] * t[0];
		//Pair<double[], double[]> result2 = trapezoidalMethod(new double[] {0, 1}, f, 1.0, 10);
		//System.out.println(Arrays.toString(result2.getSecond()));

		//nonlinearpPendulum();
		//LaneEmdenEquation();
		StiffSystem();
//		VectorFunc vf = (y, label, t) -> {
//			double[] func = new double[2];
//			func[0] = y[1] * y[1] - 2 * y[0];
//			func[1] = y[0] - y[1] - t[0] * y[1] * y[1];
//			return func;
//		};
//
//		Pair<double[], double[][]> result3 = Euler(new double[] {0, 1}, vf, new double[]{0, 1}, 10);
//		System.out.println(MVO.toString(result3.getSecond()));
		
	}

	public static void nonlinearpPendulum() {
		double g = 9.8, L = 1.0;
		double omega2 = g / L;
		VectorFunc vf = (y, label, t) -> new double[]{y[1], -omega2 * Math.sin(y[0])};
		double theta0 = 0.1 * Math.PI;
		double[] initialValue = {theta0, 0.0};
		Pair<double[], double[][]> result = trapezoidalMethod(new double[]{0.0, Math.PI}, vf, initialValue, 100);
		System.out.println(result.getSecond().length + "\t "+ result.getSecond()[0].length);
		Tecplot.LineXY("NonlinearPendulum.dat", result.getFirst(), result.getSecond());
	}

	public static void LaneEmdenEquation() {
		VectorFunc vf = (y, label, t) -> new double[]{y[1], -2 / (t[0][0] + 1.0e-6) * y[1] - Math.pow(y[0], 5.0)};
		double[] interval = {0.0, 10.0};
		double[] initValue = {1.0, 0.0};
		Pair<double[], double[][]> result = trapezoidalMethod(interval, vf, initValue, 1000);
		Tecplot.LineXY("Lane-Emden.dat", result.getFirst(), result.getSecond());
	}

	public static void StiffSystem() {
		VectorFunc vf = (y, label, t) -> new double[]{
				-10e8 * y[0] + y[1] * y[1] - 4 * y[1] + 10e8 * Math.pow(t[0][0] - 2.0, 3) + 2 * Math.pow(t[0][0] - 2, 2) + 4,
				-(1e8- 1.0) * y[0] - Math.pow(y[1], 3) + 6 * y[1] * y[1] - 12 * y[1] + 1e8 * Math.pow(t[0][0] - 2, 3) + 9
		};
		double[] interval = {0, 0.5};
		double[] initValue = {-8.0, 0.0};
		Pair<double[], double[][]> result = trapezoidalMethod(interval, vf, initValue, 100000000);
		Tecplot.LineXY("StiffSystem.dat", result.getFirst(), result.getSecond());
	}
}
